package ws_server

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/gorilla/websocket"
	"github.com/zeromicro/go-zero/core/logx"
	"net/http"
	"sync"
	"time"
)

type Server struct {
	Routes map[string]HandlerFunc
	Addr   string

	sync.RWMutex
	ConnToUser map[*WSKeepAlive]string // 连接与建立其的用户的映射
	UserToConn map[string]*WSKeepAlive // 用户与其建立连接的连接的映射

	Pattern        string         // 用于websocket通信的路由
	Authentication                // 要使用的鉴权方法
	Opt            *ServerOptions // 服务配置

	UpGrader websocket.Upgrader
	logx.Logger
}

func NewServer(addr string, fs ...ServerOptionsFunc) *Server {
	// 得到服务端的配置
	opt := newServerOptions(fs...) // f定义在options中，可用于修改服务的配置，返回一个服务配置修改后的结果

	return &Server{
		Routes: make(map[string]HandlerFunc),
		Addr:   addr,

		ConnToUser: make(map[*WSKeepAlive]string),
		UserToConn: make(map[string]*WSKeepAlive),

		Pattern:        opt.pattern,
		Authentication: opt.Authentication,
		Opt:            &opt,

		UpGrader: websocket.Upgrader{},
		Logger:   logx.WithContext(context.Background()),
	}
}

// AddRoutes 为服务建立路由和处理方法的映射，实现handler中的RegisterHandlers的逻辑
func (s *Server) AddRoutes(rs []Route) {
	for _, r := range rs {
		s.Routes[r.Method] = r.Handler
	}
}

// ServerWs 当HTTP服务收到一个与s.Pattern前缀匹配的HTTP请求时，它会调用该处理器函数处理这个请求
func (s *Server) ServerWs(w http.ResponseWriter, r *http.Request) {
	// 服务的恢复
	defer func() {
		if r := recover(); r != nil {
			s.Errorf("server handler ws recover err: %v", r)
		}
	}()

	// 将http升级为websocket
	//conn, err := s.UpGrader.Upgrade(w, r, nil)
	//if err != nil {
	//	s.Errorf("server upgrade err: %v", err)
	//	return
	//}

	// 创建自定义的wsConn（会将HTTP请求升级为websocket）
	conn := NewWSKeepAlive(s, w, r)
	if conn == nil {
		return
	}

	// 对连接鉴权
	if !s.Authentication.Auth(w, r) {
		_ = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("no privilege to access")))
		err := s.Send(&Message{FrameType: FrameData, Data: fmt.Sprint("no privilege to access")}, conn)
		if err != nil {
			s.Errorf("send message to client error: %v", err)
			return
		}
		_ = conn.Close()
		return
	}

	// 新增连接
	s.addConn(conn, r)

	// 根据连接对象进行任务处理
	go s.handlerConn(conn)

}

// handlerConn 处理连接，接受一个websocket发送的消息，根据消息的method分发到对应路由进行处理
func (s *Server) handlerConn(conn *WSKeepAlive) {
	uids := s.GetUidByConn(conn)
	conn.Uid = uids[0]

	// 协程处理已经ack的消息
	go s.handlerWrite(conn)

	// 协程处理服务端和客户端的ack确认
	if s.Opt.ackType != NoAck {
		go s.readAck(conn)
	}

	for {
		// 接受websocket发送的消息
		_, msg, err := conn.ReadMessage()
		if err != nil {
			s.Errorf("websocket conn read message err: %v", err)
			s.Close(conn)
			return
		}

		// 反序列化消息
		var message Message
		if err = json.Unmarshal(msg, &message); err != nil {
			s.Errorf("json unmarshal error: %v", err)
			s.Close(conn)
			return
		}

		if s.Opt.ackType != NoAck && message.FrameType != FrameNoAck {
			// 服务端向客户端发送ack消息
			s.Infof("conn message read ack msg %v", message)
			conn.InsertAckQueue(&message) // 如果开启了ack机制，将客户端发送过来的消息都加入AckQueue中，等待被readAck处理
		} else {
			// 无需ack的消息，将消息直接给协程任务处理（handlerWrite）
			conn.msgChan <- &message
		}

		// 根据消息类型（普通消息、心跳消息）进行处理（已改用ack机制，handlerWrite协程处理）
		//switch message.FrameType {
		//case FrameData:
		//	// 根据请求method分发路由并执行
		//	if handler, ok := s.Routes[message.Method]; ok {
		//		handler(s, conn, &message)
		//	} else {
		//		//if err = conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("method %s not found", message.Method))); err != nil {
		//		//	return
		//		//}
		//		_ = s.Send(&Message{FrameType: FrameData, Data: []byte(fmt.Sprintf("method %s not found", message.Method))}, conn)
		//		return
		//	}
		//case FramePing:
		//	_ = s.Send(&Message{FrameType: FramePing}, conn)
		//}

	}
}

// 处理服务端和客户端的ack确认
func (s *Server) readAck(conn *WSKeepAlive) {
	for {
		select {
		case <-conn.done:
			s.Infof("close message ack uid %v", conn.Uid)
			return
		default:
		}

		// 从待确认ack队列中读取最新消息
		conn.msgMu.Lock()
		if len(conn.ackQueue) == 0 {
			conn.msgMu.Unlock()
			time.Sleep(100 * time.Millisecond) // 避免了空队列时的高CPU占用
			continue
		}

		msg := conn.ackQueue[0] // 读取第一条ack消息

		switch s.Opt.ackType {
		case OnlyAck:
			// 只ack一次，给客户端发送ack消息，不需要等待客户端的对ack的ack
			_ = s.Send(&Message{
				FrameType: FrameAck,
				Id:        msg.Id,
				AckSeq:    msg.AckSeq + 1,
			}, conn)

			// 把消息从该队列中移除
			conn.ackQueue = conn.ackQueue[1:]
			conn.msgMu.Unlock()

			// 将消息传给handlerWrite处理
			conn.msgChan <- msg

		case RigorAck:
			// 需要两次ack
			// AckSeq为0，服务端确认收到，向客户端发送ack消息
			if msg.AckSeq == 0 {
				msg.AckSeq++             // 增加消息的确认序列号（会更改ackQueue中存储的值）
				msg.AckTime = time.Now() // 记录当前时间作为确认开始时间（会更改ackQueue中存储的值）
				_ = s.Send(&Message{
					FrameType: FrameAck,
					Id:        msg.Id,
					AckSeq:    msg.AckSeq,
				}, conn) // 向客户端发送ACK确认
				s.Infof("message ack RigorAck send mid %v, seq %v, time %v", msg.Id, msg.AckSeq, msg.AckTime)
				conn.msgMu.Unlock() // 解锁并继续下一轮循环，等待客户端回应
				continue
			}

			// AckSeq不为0，表明客户端收到服务端的确认，向服务端发了对确认的确认
			// 当客户端回应服务器的 ACK 确认后，会发送一个包含相同消息 ID 的响应，通过 InsertAckQueue 方法更新 msgAckRegister 映射表中对应消息的状态（AckSeq）
			newMsgSeq := conn.msgAckRegister[msg.Id].AckSeq // 从映射表中获取客户端确认消息的id，得到客户端的AckSeq

			// 如果映射表中的确认序号高于队列中的消息确认序号，表明客户端已完成对服务端ack消息的确认，服务端将消息传给handlerWrite处理并从ack队列中移除
			if newMsgSeq > msg.AckSeq {
				conn.ackQueue = conn.ackQueue[1:]
				conn.msgMu.Unlock()
				conn.msgChan <- msg
				s.Infof("message ack RigorAck success mid %v", msg.Id)
				continue
			}

			// 客户端没有确认，考虑是否超过了设置的等待时间
			// 超时，服务端结束对客户端的ack等待
			if !msg.AckTime.IsZero() && s.Opt.ackTimeout-time.Since(msg.AckTime) <= 0 {
				delete(conn.msgAckRegister, msg.Id)
				conn.ackQueue = conn.ackQueue[1:]
				conn.msgMu.Unlock()
				continue
			}

			// 未超时，服务端再次向客户端发送ack消息
			conn.msgMu.Unlock()
			_ = s.Send(&Message{
				FrameType: FrameAck,
				Id:        msg.Id,
				AckSeq:    msg.AckSeq,
			}, conn)
			time.Sleep(2 * time.Second)
		}
	}
}

// 处理根据路由websocket传递的消息
func (s *Server) handlerWrite(conn *WSKeepAlive) {
	for {
		select {
		case <-conn.done:
			return
		case message := <-conn.msgChan:
			// 根据消息的类型（普通消息、心跳消息）进行处理
			switch message.FrameType {
			case FrameData:
				// 根据请求method分发路由并执行
				if handler, ok := s.Routes[message.Method]; ok {
					handler(s, conn, message)
				} else {
					_ = s.Send(&Message{FrameType: FrameData, Data: []byte(fmt.Sprintf("method %s not found", message.Method))}, conn)
					return
				}
			case FramePing:
				_ = s.Send(&Message{FrameType: FramePing}, conn)
			}

			// 消息处理完成后，清理注册表
			if s.Opt.ackType != NoAck && message.FrameType != FrameNoAck {
				conn.msgMu.Lock()
				delete(conn.msgAckRegister, message.Id)
				conn.msgMu.Unlock()
			}
		}
	}
}

// AddConn 建立相关的映射
func (s *Server) addConn(conn *WSKeepAlive, r *http.Request) {
	// 使用server中注册的鉴权中间件的GetUserId方法，获取请求的userId
	uid := s.Authentication.GetUserId(r)

	// 对连接对象进行存储
	s.RWMutex.Lock()
	defer s.RWMutex.Unlock()

	// 目前逻辑中用户不能同时存在多个连接
	if c := s.UserToConn[uid]; c != nil {
		// 关闭之前的连接
		_ = c.Close()
	}

	// 建立连接与创建其的用户的映射
	s.UserToConn[uid] = conn
	s.ConnToUser[conn] = uid
}

func (s *Server) GetConnByUid(uid ...string) []*WSKeepAlive {
	if len(uid) == 0 {
		return nil
	}

	s.RWMutex.RLock()
	defer s.RWMutex.RUnlock()
	conns := make([]*WSKeepAlive, 0)
	for _, u := range uid {
		if conn, ok := s.UserToConn[u]; ok {
			conns = append(conns, conn)
		}
	}
	return conns
}

func (s *Server) GetUidByConn(conn ...*WSKeepAlive) []string {
	s.RLock()
	defer s.RUnlock()

	var uids []string
	if len(conn) == 0 {
		// 获取全部连接用户
		uids = make([]string, 0, len(s.ConnToUser))
		for _, c := range s.ConnToUser {
			uids = append(uids, c)
		}
	} else {
		// 获取指定连接用户
		uids = make([]string, 0, len(conn))
		for _, v := range conn {
			if c, ok := s.ConnToUser[v]; ok {
				uids = append(uids, c)
			}
		}
	}
	return uids
}

// Close 封装了WSKeepAlive连接的关闭以及连接和用户映射的删除
func (s *Server) Close(conn *WSKeepAlive) {
	_ = conn.Close() // 调用重写的连接关闭方法
	s.RWMutex.Lock()
	defer s.RWMutex.Unlock()

	uid := s.ConnToUser[conn]
	// 防止重复删除
	if uid == "" {
		return
	}
	delete(s.ConnToUser, conn)
	delete(s.UserToConn, uid)
}

// Send 向传入的连接发送消息
func (s *Server) Send(msg any, conn ...*WSKeepAlive) error {
	if len(conn) == 0 {
		return nil
	}

	data, err := json.Marshal(msg)
	if err != nil {
		return err
	}

	for _, c := range conn {
		// 调用重写的WriteMessage方法
		if err = c.WriteMessage(websocket.TextMessage, data); err != nil {
			return err
		}
	}
	return nil
}

// SendByUid 根据用户id向连接发送消息
func (s *Server) SendByUid(msg any, uid ...string) error {
	if len(uid) == 0 {
		return nil
	}

	return s.Send(msg, s.GetConnByUid(uid...)...)
}

func (s *Server) Start() {
	http.HandleFunc(s.Pattern, s.ServerWs)          // 用于将URL路径（s.Pattern）与处理HTTP请求的处理器函数（s.ServerWs）关联起来
	s.Logger.Info(http.ListenAndServe(s.Addr, nil)) // http.ListenAndServe 启动HTTP服务器
}

func (s *Server) Stop() {
	fmt.Println("Stop WebSocket ...")
}
